//      cds_vom.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include <stack>
#include <iomanip>
#include <boost/foreach.hpp>
#include <boost/math/distributions/chi_squared.hpp>

#include "cds_vom.h"

vom_node::vom_node(char c, const std::string* alphabet, 
		   vom_node* parent, int periodic)
  : c_(c), alphabet_(alphabet), 
    f_(periodic, zero_vec_t(alphabet->length())), 
    child_(alphabet->length(), NULL), parent_(parent), 
    n_occurence_(periodic, 0), lambda_(periodic, 0)
{
}

vom_node::~vom_node()
{ 
  BOOST_FOREACH(vom_node* child, child_) delete child; 
}

int vom_node::periodic() const 
{ 
  return sz(f_); 
}

bool vom_node::has_table(int pr) const 
{ 
  return sz(f_[pr]) == alphabet_size();
}

std::string vom_node::str() const
{
  const vom_node* N = this;
  std::string s;
	
  while (N->parent_ != NULL) {
    s = s + N->c_;
    N = N->parent_;
  }
  return s;
}

int vom_node::alphabet_size() const
{
  return sz(*alphabet_);
}

char vom_node::alphabet(int j) const
{
  return alphabet_->at(j);
}

void vom_node::normalize()
{
  for (int frame = 0; frame < periodic(); frame++) {
    //BOOST_FOREACH(vec_t& ft, f_) {
    vec_t& ft = f_[frame];
    //int n = n_occurence_[frame];

    double min = DOUBLE_INFINITY;
    bool has_zero = false;
    double smooth = 0.1;
    for (int i = 0; i < sz(ft); i++) {
      if (min > ft[i] && ft[i] > 1e-16) min = ft[i];
      if (ft[i] == 0) has_zero = true;
    }
    if (min == DOUBLE_INFINITY) min = 1;
    
    if (has_zero) ft += scalar_vec_t(sz(ft), min * smooth);
    ft /= ublas::sum(ft);
  }
}

void vom_node::set_frequency(const sastring_vec_t& train_set)
{
  std::string s = this->str();
  int n = sz(s);

  BOOST_FOREACH(vec_t& ft, f_) ft = zero_vec_t(sz(ft));
  n_occurence_ = vi_t(periodic(), 0);

  bool has_to_search = true;
  if (parent_ != NULL) {
    int sum = 0;
    BOOST_FOREACH(int n, parent_->n_occurence_) sum += n;
    if (sum == 0) has_to_search = false;
  }

  if (has_to_search) {
    BOOST_FOREACH(const sastring& train, train_set) {
      int left, size;
      train.sasearch(s, left, size);
      
      for (int i = 0; i < size; i++) { // all occurences of context s
	int pos = train.SA[left+i];
	if (pos + n >= sz(train)) continue; // end of string
	
	char c = train[pos+n];
	int frame = (pos+n) % periodic();
	
	n_occurence_[frame]++;
	for (int j = 0; j < alphabet_size(); j++)
	  if (alphabet(j) == c) f_[frame](j) += 1.0;
      }
    }
  }
  normalize();
}

void vom_node::copy(vom_node* N)
{
  // copy normal field
  *this = *N;
  // copy pointer field
  for (int i = 0; i < alphabet_size(); i++) {
    char c = alphabet(i);
    if (N->child_[i] != NULL) { // if child is not NULL
      child_[i] = new vom_node(c, alphabet_, this, periodic()); 
      child_[i]->copy(N->child_[i]);
      child_[i]->parent_ = this;
    }
  }
}

vec_t vom_node::prob_imm(int frame)
{
  vom_node* N = this;
  double lambda = 1;
  vec_t f = zero_vec_t(alphabet_size());
  while (N != NULL) {
    f += lambda * N->lambda_[frame] * N->f_[frame];
    lambda *= 1-N->lambda_[frame];
    N = N->parent_;
  }
  return f;
}

// ====================== VOM =============================
VOM::VOM(const std::string& alpha, int periodic) 
  : alphabet(alpha), root_(NULL), eta_(0.1), T_(100), periodic_(periodic),
    count_p_(periodic,0), depth_p_(periodic,0) 
{
}

VOM::~VOM() 
{ 
  delete root_; 
}

void VOM::copy(VOM& m)
{
  // copy normal field
  *this = m;
  // pointer field
  root_ = new vom_node('X', &alphabet, NULL, periodic_);
  root_->copy(m.root_);
}

void VOM::print_node(std::ostream& os) const 
{ 
  print_node(os, root_, 0); 
}

void VOM::print_node(std::ostream& os, vom_node* N, int depth) const
{
  if (N == NULL) {
    os << "UP" << ENDL;
    return;
  }
  //for (int i = 0; i < depth; i++) os << "  ";

  if (N->parent_ == NULL) os << "X";
  else os << N->str();
  BOOST_FOREACH(vec_t& ft, N->f_) {
    if (sz(ft) == sz(alphabet))
      for (int i = 0; i < sz(ft); i++) os << " " << ft[i];
    else
      for (int i = 0; i < sz(alphabet); i++) os << " " << 0;
  }
  os << ENDL;
	
  BOOST_FOREACH(vom_node* child, N->child_)
    print_node(os, child, depth+1);
}

void VOM::load_node(std::istream& is) 
{ 
  count_ = 0; 
  load_node(is, &root_, NULL, 0); 
}

void VOM::load_node(std::istream& is, vom_node** N, vom_node* parent, int depth)
{
  FileAccessLineSplit(is, line, tokens, "\r\n\t ,");
  if (sz(tokens) == 0 || tokens[0] == "UP") {
    *N = NULL;
    return;
  }
  count_ ++;
	
  (*N) = new vom_node(tokens[0][0], &alphabet, parent, periodic_);
  for (int i = 0; i < sz(alphabet)*periodic_; i++) {
    double f = s2<double>(tokens[i+1]);
    int frame = i/sz(alphabet);
    int c = i % sz(alphabet);
    (*N)->f_[frame][c] = f;
  }
  for (int fr = 0; fr < periodic_; fr++)
    if (ublas::sum((*N)->f_[fr]) == 0) (*N)->f_[fr] = vec_t();
	
  for (int i = 0; i < sz(alphabet); i++) {
    vom_node** child = &((*N)->child_[i]);
    load_node(is, child, *N, depth+1);
  }
}
		
// void VOM::grow(const sastring_vec_t& train_set, double eta) 
// { 
//   eta_ = eta; count_ = 0; grow_node(root_, train_set); 
// }

// void VOM::prune(const sastring_vec_t& train_set, double eta) 
// { 
//   eta_ = eta; 
//   prune_node(&root_, train_set); 
// }

void VOM::depth() 
{ 
  count_ = 0; depth_ = 0;
  std::fill(all_element(count_p_), 0);
  std::fill(all_element(depth_p_), 0);
  depth_node(root_, 0); 
}

void VOM::depth_node(vom_node* N, int depth)
{
  if (N == NULL) return;
  count_++;
  if (depth_ < depth) depth_ = depth;
  for (int fr = 0; fr < periodic_; fr++) {
    if (N->has_table(fr)) {
      count_p_[fr]++;
      if (depth_p_[fr] < depth) depth_p_[fr] = depth;
    }
  }
  
  BOOST_FOREACH(vom_node* child, N->child_) 
    depth_node(child, depth+1);
}

void VOM::stats_print(const std::string& statsFile) const
{
  FileAccessWrite(out, statsFile);
  out << " " << std::setw(10) << count_ 
      << " " << std::setw(10) << depth_;
  for (int i = 0; i < periodic_; i++)
    out << " " << std::setw(10) << count_p_[i] 
	<< " " << std::setw(10) << depth_p_[i]; 
  out << ENDL;
}

const vec_t& VOM::f_table(const std::string& s, int pos) const
{
  vom_node* N = find_node(s, pos);
  int frame = pos % periodic_;
  return N->f_[frame];
}

vom_node* VOM::find_node(const std::string& s, int pos) const
{
  vom_node* N = root_;
  int i = pos-1;
  bool stop = false;
  int frame = pos % periodic_;

  while (i >= 0 && !stop) { // trace upstream for longest context
    char c = s[i];
    stop = true;
    int j = this->pos(c);		
		
    // check (valid nucleotide && has child && child has probability table)
    if ( j >= 0 && N->child_[j] != NULL && N->child_[j]->has_table(frame) ) {
      stop = false;
      N = N->child_[j];
    }
    else break;
    i--;
  }

  return N;
}

double VOM::log_prob(const std::string& s) const
{
  double l = 0;
  for (int i = 0; i < sz(s); i++)
    l += log_prob(s, i);
  return l;
}

double VOM::log_prob(const std::string& s, int pos) const
{
  const vec_t& ft = f_table(s, pos);
  char c = s[pos];
  int j = this->pos(c);
  return log(ft(j));
}

int VOM::pos(char c) const
{
  for (int j = 0; j < sz(alphabet); j++) 
    if (alphabet[j] == c) return j;
  return -1;
}

char VOM::ch(int pos) const
{
  return alphabet[pos];
}

// =============== FOM builder ==================
void VOM::build_fom(const sastring_vec_t& train_set, int order)
{
  total_node_ = (pow(sz(alphabet), order+1) - 1) / (sz(alphabet)-1);
  MSG("total=" << total_node_);
  current_progress_ = 0;
  root_ = new vom_node('X', &alphabet, NULL, periodic_);
  build_fom_node(root_, train_set, order);
  MSG("");
}

void VOM::build_fom_node(vom_node* N, 
			 const sastring_vec_t& train_set, int order)
{
  if (floor(current_progress_/total_node_*100) < floor((current_progress_+1)/total_node_*100) 
      || current_progress_ == 0) {
    std::cerr << floor((current_progress_+1)/total_node_*100) << "\% ";
    std::cerr.flush();
  }
  N->set_frequency(train_set);
  current_progress_++;
  if (order > 0) {
    for (int i = 0; i < sz(alphabet); i++) {
      char c = alphabet[i];
      vom_node* child = new vom_node(c, &alphabet, N, periodic_);
      N->child_[i] = child;
    }
    BOOST_FOREACH(vom_node* child, N->child_)
      build_fom_node(child, train_set, order-1);
  }
}

void VOM::build_fom_flat(const sastring_vec_t& train_set, int order) {
  // build tree
  std::stack<vom_node*> stack;
  root_ = new vom_node('X', &alphabet, NULL, periodic_); root_->depth_ = 0;
  stack.push(root_);
  while (!stack.empty()) {
    vom_node* N = stack.top(); stack.pop();
    if (N->depth_ >= order) continue;
    for (int i = 0; i < sz(alphabet); i++) {
      char c = alphabet[i];
      vom_node* child = new vom_node(c, &alphabet, N, periodic_); 
      child->depth_ = N->depth_ + 1;
      stack.push(child);
      N->child_[i] = child;
    }
  }
  // accumulate occurences
  BOOST_FOREACH(const std::string& s, train_set) {
    for (int i = 0; i < sz(s); i++) accumulate_occurence(s, i, order);
  }
  // normalize
  stack.push(root_);
  while (!stack.empty()) {
    vom_node* N = stack.top(); stack.pop();
    N->normalize();
    BOOST_FOREACH(vom_node* child, N->child_) 
      if (child != NULL) stack.push(child);
  }
}

void VOM::accumulate_occurence(const std::string& s, int p, int order) 
{
  int frame = p % periodic_;
  char b = s[p];
  int j_b = pos(b);
  vom_node* N = root_;
  N->f_[frame](j_b) += 1.0;
  N->n_occurence_[frame] += 1;

  for (int i = p-1; i >= 0 && i >= p-order; i--) {
    char c = s[i];
    int j_c = pos(c);
    N = N->child_[j_c];
    N->f_[frame](j_b) += 1.0;
    N->n_occurence_[frame] += 1;
  }
}


// ========================= IMM builder ============================
void VOM::compute_lambda_xi_square()
{
  std::vector<vec_t> imm_f = std::vector<vec_t>(periodic_, zero_vec_t(sz(alphabet)));
  compute_lambda_xi_square(root_, imm_f);
}

void VOM::compute_lambda_deleted_interpolation()
{
}

double VOM::log_prob_lambda(const std::string& s) const
{
  double l = 0;
  for (int pos = 0; pos < sz(s); pos++) { 
    char c_pos = s[pos];
    int j_pos = this->pos(c_pos);
    int frame = pos % periodic_;

    vom_node* N = root_;
    double p = N->f_[frame](j_pos);
    int i = pos-1;
    bool stop = false;
    while (i >= 0 && !stop) { // trace upstream and interpolate using lambda
      char c = s[i];
      stop = true;
      int j = this->pos(c);
      if (j >= 0 && N->child_[j] != NULL && N->child_[j]->has_table(frame) ) {
	stop = false;
	N = N->child_[j];
	double lambda = N->lambda_[frame];
	if (lambda == 0) break;
	p = lambda * N->f_[frame](j_pos) + (1-lambda) * p;
      }
      else break;
      i--;
    }
    
    l += log(p);
  }
  return l;
}

void VOM::compute_lambda_xi_square(vom_node* N, const std::vector<vec_t>& imm_f)
{
  if (N == NULL) return;
  int T = T_;
  // compute lambda
  boost::math::chi_squared chi(3); // 3 degree of freedom

  std::vector<vec_t> new_imm_f = std::vector<vec_t>(periodic_, zero_vec_t(sz(alphabet)));
  if (N->parent_ == NULL) N->lambda_ = scalar_vec_t(periodic_, 1.0);
  else for (int frame = 0; frame < periodic_; frame++) { // for each frame
      double& lambda = N->lambda_[frame];
      const vec_t& imm_ft = imm_f[frame];
      const vec_t& ft = N->f_[frame];
      
      if (N->n_occurence_[frame] > T) lambda = 1;
      else {
	double xi = 0;
	for (int j = 0; j < sz(alphabet); j++) 
	  xi += sqr(ft[j]-imm_ft(j)) / imm_ft(j);
	double q = boost::math::cdf(chi, xi); // 1 minus p-value
	lambda = (q < 0.5) ? 0 : q * N->n_occurence_[frame] / T;
	//if (q<0.5) MSG("less than 0.5"); else MSG("greater than 0.5");
      }

      // new_imm_f
      new_imm_f[frame] = lambda*N->f_[frame] + (1-lambda)*imm_f[frame];
    }

  // recursively on children
  BOOST_FOREACH(vom_node* c, N->child_) compute_lambda_xi_square(c, new_imm_f);
}

// ========================= prune ==================================
void VOM::prune(const sastring_vec_t& train_set, double eta)
{
  eta_ = eta; 
  prune_node(&root_, train_set);
}

void VOM::prune_node(vom_node** N, const sastring_vec_t& train_set)
{  
  if (*N == NULL) return;
  for (int i = 0; i < sz(alphabet); i++) {
    vom_node** child = &((*N)->child_[i]);
    prune_node(child, train_set);
  }
	
  vom_node* parent = (*N)->parent_;
  if (parent == NULL) return; // cannot prune root_

  for (int pr = 0; pr < periodic_; pr++) {
    // if at least one child of N has this periodic table then we cannot prune
    bool ok = true;
    for (int i = 0; i < sz(alphabet); i++) 
      if ((*N)->child_[i] != NULL && (*N)->child_[i]->has_table(pr)) ok = false;
    if (!ok) continue;

    // otherwise, prune this periodic table if KL distance is small
    if ((*N)->has_table(pr) && parent->has_table(pr)) {
      double KL = CommonFunction::KLdivergence((*N)->f_[pr], parent->f_[pr]);
      if (KL <= eta_) (*N)->f_[pr] = vec_t();
    }
  }
   
  // if all periodic tables are pruned then prune the node
  bool ok = false;
  for (int pr = 0; pr < periodic_; pr++) 
    if ((*N)->has_table(pr)) ok = true; 
  if (!ok) {
    delete *N;
    *N = NULL;
  }
}

void VOM::prune_lambda(const sastring_vec_t& train_set, double eta)
{
  eta_ = eta; 
  prune_lambda_node(&root_, train_set);
}

void VOM::prune_lambda_node(vom_node** N, const sastring_vec_t& train_set)
{  
  if (*N == NULL) return;
  for (int i = 0; i < sz(alphabet); i++) {
    vom_node** child = &((*N)->child_[i]);
    prune_node(child, train_set);
  }
	
  vom_node* parent = (*N)->parent_;
  if (parent == NULL) return; // cannot prune root_

  for (int pr = 0; pr < periodic_; pr++) {
    // if at least one child of N has this periodic table and lambda > 0 then we cannot prune
    bool ok = true;
    for (int i = 0; i < sz(alphabet); i++) 
      if ((*N)->child_[i] != NULL && (*N)->child_[i]->has_table(pr) 
	  && (*N)->child_[i]->lambda_[pr] > 0) ok = false;
    if (!ok) continue;

    // otherwise, prune this periodic table if KL distance is small
    if ((*N)->has_table(pr) && parent->has_table(pr)) {
      vec_t p_imm = parent->prob_imm(pr);
      const vec_t& p = (*N)->f_[pr];
      const double& lambda = (*N)->lambda_[pr];
      double KL = 0; //CommonFunction::KLdivergence((*N)->f_[pr], parent->f_[pr]);
      for (int i = 0; i < sz(alphabet); i++) {
	KL += p(i) * log(lambda * p(i) / p_imm(i) + 1-lambda);
      }
      if (KL <= eta_ || lambda == 0) {
	(*N)->f_[pr] = vec_t();
	(*N)->lambda_[pr] = 0;
      }
    }
  }
   
  // if all periodic tables are pruned then prune the node
  bool ok = false;
  for (int pr = 0; pr < periodic_; pr++) 
    if ((*N)->has_table(pr)) ok = true; 
  if (!ok) {
    delete *N;
    *N = NULL;
  }
}

